﻿/*
 * PSMuxer.cpp
 *
 *  Created on: 2016年2月26日
 *      Author: terry
 */

#include "PSMuxer.h"
#include "Ffmpeg.h"
#include "put_bits.h"
#include "mpeg.h"
#include "H264PropParser.h"



namespace av
{

static uint8_t s_IMKH[] =
{
	0x49, 0x4D, 0x4B, 0x48, 0x01, 0x01, 0x00, 0x00,  0x02, 0x00, 0x00, 0x01, 0x11, 0x71, 0x01, 0x10,
	0x40, 0x1F, 0x00, 0x00, 0x00, 0xFA, 0x00, 0x00,  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
};

static int s_mux_rate = 0;

#define MAX_PES_PACKET_LENGTH  0xFFF0


int PSMuxer::put_pack_header(uint8_t *buf, int64_t timestamp)
{
    PutBitContext pb;

    init_put_bits(&pb, buf, 128);

    put_bits32(&pb, PACK_START_CODE);
    put_bits(&pb, 2, 0x1);
    put_bits(&pb,  3, (uint32_t)((timestamp >> 30) & 0x07));
    put_bits(&pb,  1, 1);
    put_bits(&pb, 15, (uint32_t)((timestamp >> 15) & 0x7fff));
    put_bits(&pb,  1, 1);
    put_bits(&pb, 15, (uint32_t)((timestamp)       & 0x7fff));
    put_bits(&pb,  1, 1);

    uint32_t scrExt = 0x10B;
    /* clock extension */
    put_bits(&pb, 9, scrExt);

    put_bits(&pb, 1, 1);

    put_bits(&pb, 22, s_mux_rate);
    put_bits(&pb, 1, 1);
    put_bits(&pb, 1, 1);
    put_bits(&pb, 5, 0x1f); /* reserved */
    put_bits(&pb, 3, 6); /* stuffing length */

	time_t t = time(NULL);
	struct tm tmNow = *localtime(&t);
	int year = (tmNow.tm_year + 1900) - 2000;
	int month = tmNow.tm_mon + 1;
	int day = tmNow.tm_mday;

	put_bits(&pb, 8, (uint8_t)year);
	put_bits(&pb, 8, (uint8_t)month);
	put_bits(&pb, 8, (uint8_t)day);
	put_bits(&pb, 8, (uint8_t)tmNow.tm_hour);
	put_bits(&pb, 8, (uint8_t)tmNow.tm_min);
	put_bits(&pb, 8, (uint8_t)tmNow.tm_sec);

    flush_put_bits(&pb);
    return put_bits_ptr(&pb) - pb.buf;
}

int PSMuxer::put_system_header(uint8_t *buf, int stream_id)
{
    PutBitContext pb;

    init_put_bits(&pb, buf, 128);

    put_bits32(&pb, SYSTEM_HEADER_START_CODE);
    put_bits(&pb, 16, 0);
    put_bits(&pb, 1, 1);

    /* maximum bit rate of the multiplexed stream */
    put_bits(&pb, 22, s_mux_rate);
    put_bits(&pb, 1, 1); /* marker */

    if ((stream_id & 0xE0) == VIDEO_ID) {
        /* This header applies only to the video stream
         * (see VCD standard p. IV-7) */
        put_bits(&pb, 6, 0);
    } else
        put_bits(&pb, 6, 0x3f); // audio bound ?

    put_bits(&pb, 1, 0); /* variable bitrate */
    put_bits(&pb, 1, 0); /* non constrainted bit stream */

    put_bits(&pb, 1, 0); /* audio locked */
    put_bits(&pb, 1, 0); /* video locked */

    put_bits(&pb, 1, 1); /* marker */

    if ((stream_id & 0xe0) == AUDIO_ID) {
        /* This header applies only to the audio stream
         * (see VCD standard p. IV-7) */
        put_bits(&pb, 5, 0);
    } else
        put_bits(&pb, 5, 1); // video bound ?

    put_bits(&pb, 1, 0);    /* packet_rate_restriction_flag */
    put_bits(&pb, 7, 0x7f); /* reserved byte */
    // put_bits(&pb, 8, 0xff); /* reserved byte */ ? 对比

    put_bits(&pb, 8, stream_id);         /* stream ID */
    put_bits(&pb, 2, 3);
    if (stream_id < 0xe0) {
        /* audio */
        put_bits(&pb, 1, 0);
        put_bits(&pb, 13, 232); //? max_buffer_size?
    } else {
        /* video */
        put_bits(&pb, 1, 1);
        put_bits(&pb, 13, 32); //?
    }

    flush_put_bits(&pb);
    int size = put_bits_ptr(&pb) - pb.buf;
    /* patch packet size */
    AV_WB16(buf + 4, size - 6);

    return size;
}

uint8_t PSMuxer::get_stream_type(int stream_codec_id)
{
    uint8_t res = 0;
    switch(stream_codec_id)
    {
    case AV_CODEC_ID_PCM_ALAW:
        res = 0x90;
        break;
    case AV_CODEC_ID_PCM_MULAW:
        res = 0x91;//STREAM_TYPE_AUDIO_PCM_MULAW;
        break;
    case AV_CODEC_ID_MPEG2VIDEO:
        res = STREAM_TYPE_VIDEO_MPEG2;
        break;
    case AV_CODEC_ID_MP3:
        res = STREAM_TYPE_AUDIO_MPEG2;
        break;
    case AV_CODEC_ID_MPEG4:
        res = STREAM_TYPE_VIDEO_MPEG4;
        break;
    case AV_CODEC_ID_AC3:
        res = STREAM_TYPE_AUDIO_AC3;
        break;
    case AV_CODEC_ID_AAC:
        res = STREAM_TYPE_AUDIO_AAC;
        break;
    case AV_CODEC_ID_H264:
        res = STREAM_TYPE_VIDEO_H264;
        break;
	case AV_CODEC_ID_HEVC:
		res = STREAM_TYPE_VIDEO_HEVC;
		break;
    default:
        break;
    }
    return res;
}


int PSMuxer::put_psm_header(uint8_t *buf, int videoId, int audioId)
{
    PutBitContext pb;

    int streamnum = 0;
    if (videoId != AV_CODEC_ID_NONE)
    {
        streamnum ++;
    }
    if (audioId != AV_CODEC_ID_NONE)
    {
        streamnum ++;
    }

    int streminfolen = streamnum * 4;

    init_put_bits(&pb, buf, (16 + streminfolen) * 8);

    put_bits32(&pb, 0x000001bc);
    put_bits(&pb, 16, 10 + streminfolen);         /*program stream map length*/
    put_bits(&pb, 1, 1);          /*current next indicator */
    put_bits(&pb, 2, 3);          /*reserved*/
    put_bits(&pb, 5, 0);          /*program stream map version*/
    put_bits(&pb, 7, 0x7F);       /*reserved */
    put_bits(&pb, 1, 1);          /*marker bit */
    put_bits(&pb, 16,0);          /*programe stream info length*/
    put_bits(&pb, 16, streminfolen);         /*elementary stream map length  is*/

    if (videoId != AV_CODEC_ID_NONE)
    {
        put_bits(&pb, 8, get_stream_type(videoId));       /*stream_type*/
        put_bits(&pb, 8, VIDEO_ID + 1);       /*elementary_stream_id*/
        put_bits(&pb, 16, 0);         /*elementary_stream_info_length is*/
    }

    if (audioId != AV_CODEC_ID_NONE)
    {
        put_bits(&pb, 8, get_stream_type(audioId));       /*stream_type*/
        put_bits(&pb, 8, AUDIO_ID + 1);       /*elementary_stream_id*/
        put_bits(&pb, 16, 0);         /*elementary_stream_info_length is*/
    }

    uint32_t crc = 0;//crc32(0, buf, 13 + streminfolen);
    /*audio*/
    /*crc (2e b9 0f 3d)*/
    put_bits32(&pb, crc);       /*crc */
    flush_put_bits(&pb);

    return 16 + streminfolen;
}

int PSMuxer::put_pes_header(uint8_t *buf, int stream_id, int pktSize, uint32_t ptsDtsFlag, int64_t pts, int64_t dts)
{
    uint8_t lPesHeaderLen = (ptsDtsFlag == 2) ? 0x07 : (0x0A+2);
    int pesLength = 3 + lPesHeaderLen + pktSize - 2;

    PutBitContext pb;
    init_put_bits(&pb, buf, 128);

    put_bits(&pb, 24, 0x000001);    //startcode
    put_bits(&pb, 8, stream_id);
    put_bits(&pb, 16, pesLength);

    put_bits(&pb, 8, 0x80);
    put_bits(&pb, 2, ptsDtsFlag);
    put_bits(&pb, 6, 0);			//flags

    if (ptsDtsFlag == 2)
    {
        put_bits(&pb, 8, 0x05);     //pes_header_data_length;
        put_bits(&pb, 4, 2);
        put_bits(&pb, 3, (int)(pts>>30));
        put_bits(&pb, 1, 1);
        put_bits(&pb, 15, (int)((pts&0x000000003fff8000ll)>>15));
        put_bits(&pb, 1, 1);
        put_bits(&pb, 15, (int)(pts&0x0000000000007fffll));
        put_bits(&pb, 1, 1);

    }
    else
    {
        put_bits(&pb, 8, 0x0A);     //pes_header_data_length;
        put_bits(&pb, 4, 3);
        put_bits(&pb, 3, (int)(pts>>30));
        put_bits(&pb, 1, 1);
        put_bits(&pb, 15, (int)((pts&0x000000003fff8000ll)>>15));
        put_bits(&pb, 1, 1);
        put_bits(&pb, 15, (int)(pts&0x0000000000007fffll));
        put_bits(&pb, 1, 1);

        put_bits(&pb, 4, 1);
        put_bits(&pb, 3, (int)(dts>>30));
        put_bits(&pb, 1, 1);
        put_bits(&pb, 15, (int)((dts&0x000000003fff8000ll)>>15));
        put_bits(&pb, 1, 1);
        put_bits(&pb, 15, (int)(dts&0x0000000000007fffll));
        put_bits(&pb, 1, 1);

    }

    flush_put_bits(&pb);

    int size = put_bits_ptr(&pb) - pb.buf;
    return size;
}

PSMuxer::PSMuxer():
		m_vCodec(),
		m_aCodec()
{
}

PSMuxer::~PSMuxer()
{
}

bool PSMuxer::setCodec(int vCodec, int aCodec)
{
	if ((vCodec <= 0) && (aCodec <= 0))
	{
		return false;
	}

	m_vCodec = vCodec;
	m_aCodec = aCodec;
	return true;
}

bool PSMuxer::writeVideo(uint8_t *pktData, int pktSize, bool keyFrame, int64_t pts, uint8_t *& outData, int& outSize)
{
	if ((!pktData) || (pktSize <= 0))
	{
		return false;
	}

	m_buffer.clear();

	int streamID = VIDEO_ID + 1;
	int size = 0;

	if (keyFrame)
	{
		size += writePackHeader(pts);
		size += writeSystemHeader(streamID);

		size += writePsmHeader(m_vCodec, m_aCodec);

		size_t idx = H264PropParser::findStartCode(pktData, pktSize, 4);
		if (idx != -1)
		{
			size += writePesPacket(streamID, pktData, idx, 3, pts, pts);

			pktData += idx;
			pktSize -= idx;

			idx = H264PropParser::findStartCode(pktData, pktSize, 4);
			while (idx != -1)
			{
				size += writePesPacket(streamID, pktData, idx, 3, pts, pts);

				pktData += idx;
				pktSize -= idx;

				idx = H264PropParser::findStartCode(pktData, pktSize, 4);
			}
		}

		size += writePesPacket(streamID, pktData, pktSize, 3, pts, pts);
	}
	else
	{
		size += writePesPacket(streamID, pktData, pktSize, 3, pts, pts);
	}

	outData = m_buffer.data();
	outSize = size;

	return true;
}

bool PSMuxer::writeAudio(uint8_t *pktData, int pktSize, int64_t pts, uint8_t *& outData, int& outSize)
{
	if (m_aCodec <= 0)
	{
		return false;
	}

	if ((!pktData) || (pktSize <= 0))
	{
		return false;
	}

	m_buffer.clear();

	int streamID = AUDIO_ID + 1;
	int size = 0;

	//size += writePackHeader(pts); /// 海康流的音频流前面没有这个包

	size += writePesPacket(streamID, pktData, pktSize, 3, pts, pts);

	outData = m_buffer.data();
	outSize = size;

	return true;
}

int PSMuxer::writePackHeader(int64_t pts)
{
    uint8_t buf[128] = { 0 };
    int size = put_pack_header(buf, pts);
    m_buffer.write(buf, size);
    return size;
}

int PSMuxer::writeSystemHeader(int streamID)
{
    uint8_t buf[128] = { 0 };
    int size = put_system_header(buf, streamID);
    m_buffer.write(buf, size);
    return size;
}

int PSMuxer::writePsmHeader(int videoCodec, int audioCodec)
{
    uint8_t buf[128] = {0};
    int size = put_psm_header(buf, videoCodec, audioCodec);
    m_buffer.write(buf, size);
    return size;
}

int PSMuxer::writePesPacket(int streamID, uint8_t *data, int size, int flags, int64_t pts, int64_t dts)
{
    int totalSize = 0;
    while (size > 0)
    {
        int pktSize = (size > MAX_PES_PACKET_LENGTH) ? MAX_PES_PACKET_LENGTH : size;

        uint8_t hdrBuf[128] = {0};
        int hdrSize = put_pes_header(hdrBuf, streamID, pktSize, flags, pts, dts);

        m_buffer.write(hdrBuf, hdrSize);
        m_buffer.write(data, pktSize);

        data += pktSize;
        size -= pktSize;

        totalSize += hdrSize + pktSize;

    }
    return totalSize;
}


bool PSMuxer::writeIMKH(uint8_t *& outData, int& outSize)
{
	outData = s_IMKH;
	outSize = sizeof(s_IMKH);
	return true;
}

void PSMuxer::clear()
{
	m_buffer.clear();
}


} /* namespace av */
